
#include "scatter_reduce_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include <algorithm>

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

  ScatterReduceTilingData tiling;
  auto self_shape = context->GetInputShape(0)->GetOriginShape();
  auto index_shape = context->GetInputShape(1)->GetOriginShape();
  auto src_shape = context->GetInputShape(2)->GetOriginShape();

    int index_dim = index_shape.GetDimNum();

    int32_t index_size = 1;
    int32_t self_size = 1;
    int32_t src_size = 1;
    int32_t self_ndarray[8];
    int32_t index_ndarray[8];
    int32_t src_ndarray[8];

    for (int i = 0; i < index_dim; i++)
    {
        self_ndarray[i] = self_shape.GetDim(i);
        index_ndarray[i] = index_shape.GetDim(i);
        src_ndarray[i] = src_shape.GetDim(i);
        index_size *= index_ndarray[i];
        self_size *= self_ndarray[i];
        src_size *= src_ndarray[i];
    }
    int dim = *context->GetAttrs()->GetInt(0);
    if(dim < 0) dim += index_dim;
    const char *reduce_attr = context->GetAttrs()->GetAttrPointer<char>(1);
    bool include_self = *context->GetAttrs()->GetBool(2);

    uint8_t reduce = 0;
    if(strcmp(reduce_attr, "sum") == 0) {
        reduce = 0;
    }
    else if(strcmp(reduce_attr, "prod") == 0) {
        reduce = 1;
    }
    else if(strcmp(reduce_attr, "mean") == 0) {
        reduce = 2;
    }
    else if(strcmp(reduce_attr, "amax") == 0) {
        reduce = 3;
    }
    else  {//amin
        reduce = 4;
    }
    tiling.set_reduce(reduce);
    tiling.set_include_self(include_self);

    int32_t loopCount = 1;
    int32_t cycles = 1;
    int32_t interval = 1;

    for(int i = 0; i < dim; i++)
    {
        loopCount *= index_ndarray[i];
    }
    
    cycles = index_ndarray[dim];
    
    for(int i = dim+1; i < index_dim; i++)
    {
        interval *= index_ndarray[i];
    }

    uint32_t inputBytes = GetSizeByDataType(context->GetInputDesc(0)->GetDataType()); //输入类型

    if(index_size == self_size && index_size == src_size && reduce == 4 && include_self == false && interval >= 256 && ((interval*inputBytes) % 64 == 0))
    {
        int8_t formerNum;
        int16_t tailLength;
        auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
        auto aivNum = ascendcPlatform.GetCoreNumAiv();

        aivNum = 40/2;

        if(cycles < aivNum)
        {
            formerNum = cycles;
            tailLength = 0;
            context->SetBlockDim(formerNum*2);
        }
        else
        {
            formerNum = cycles % aivNum;
            tailLength = cycles / aivNum;
            if(formerNum == 0) 
            {
                formerNum = aivNum;
                tailLength -= 1;
            }
            context->SetBlockDim(aivNum*2);
        }
        
        tiling.set_loopCount(loopCount);
        tiling.set_cycles(cycles);
        tiling.set_interval(interval/2);

        tiling.set_formerNum(formerNum);
        tiling.set_tailLength(tailLength);
        
        context->SetTilingKey(1);
        
    }
    else
    {
        tiling.set_index_dim(index_dim);
        tiling.set_dim(dim);

        tiling.set_self_size(self_size);
        tiling.set_index_size(index_size);
        tiling.set_self_ndarray(self_ndarray);
        tiling.set_index_ndarray(index_ndarray);
        tiling.set_src_ndarray(src_ndarray);

        context->SetTilingKey(2);
        context->SetBlockDim(1);
    }

    
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class ScatterReduce : public OpDef {
public:
    explicit ScatterReduce(const char* name) : OpDef(name)
    {
        this->Input("self")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("index")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("src")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("dim").Int();
        this->Attr("reduce").String();
        this->Attr("include_self").AttrType(OPTIONAL).Bool(true);

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");

    }
};

OP_ADD(ScatterReduce);
}
